# -*- coding: utf-8 -*-

import os
import cv2
import paddle
import numpy as np
from src.model import Generator
from src import parameter

#基本参数
paraPath = "src/work/generator.params"
outputPath = "output/imgs"
if not os.path.exists(outputPath):
    os.makedirs(outputPath)
    

#生成器网络
netG = Generator()
#加载示例参数
layer_state_dict = paddle.load(paraPath)
netG.set_state_dict(layer_state_dict)

count = 0
while True:
    noise = paddle.randn([1,parameter.G_DIMENSION,1,1],'float32')
    generated_image = netG(noise).numpy() 
    image = generated_image[0].transpose()
    image = np.where(image > 0, image, 0)
    image = np.where(image < 1, image, 1)
    image = image*255.
    image = image.astype(np.uint8)
    
    image = cv2.resize(image,(256,256))
    cv2.imshow("generator",image)
    key = cv2.waitKey(0)
    if key & 0xff == ord("q"):
        break
    elif key & 0xff == ord("s"):
        cv2.imwrite(os.path.join(outputPath,"{}.jpg".format(count)),image)
        count+=1
    else:
        pass

cv2.destroyAllWindows()












